
#ifndef RTCONNECTORSOCKSPROXYT_H
#define RTCONNECTORSOCKSPROXYT_H

#include "RtConnectorTcpT.h"
#include "RtTransportTcp.h"
#include "RtHttpUtilClasses.h"
#include "RtMessageBlock.h"

template <class UpperType, class UpTrptType, class UpSockType>
class /*RT_API_EXPORT*/ CRtConnectorSocksProxyT 
	: public IRtConnectorInternal
	, public IRtTransportSink 
{
public:
	typedef CRtConnectorSocksProxyT SelfType;

	CRtConnectorSocksProxyT(
			IRtReactor *aReactor, 
			UpperType &aUpper, 
			CRtHttpProxyInfo* aProxyInfo = NULL)
		: m_pReactor(aReactor)
		, m_Upper(aUpper)
		, m_TcpConnector(aReactor, *this)
		, m_pTransport(NULL)
		, m_State(STATE_IDLE)
		, m_dwSvrIP(0)
		, m_wSvrPort(0)
	{
		if (aProxyInfo)
			SetProxyInfo(aProxyInfo);
	}
	
	virtual ~CRtConnectorSocksProxyT()
	{
		Close();
	}

	void SetProxyInfo(CRtHttpProxyInfo* aProxyInfo)
	{
		RT_ASSERTE_RETURN_VOID(aProxyInfo);

		m_pProxyInfo = aProxyInfo;
		RT_ASSERTE(
			m_pProxyInfo->GetProxyType() == CRtHttpProxyInfo::SOCK4_PROXY || 
			m_pProxyInfo->GetProxyType() == CRtHttpProxyInfo::SOCK5_PROXY);
	}

	// interface IRtConnectorInternal
	virtual int Connect(const CRtInetAddr &aAddr, CRtInetAddr *aAddrLocal = NULL)
	{
		RT_ASSERTE(m_State == STATE_IDLE);
		RT_ASSERTE_RETURN(!m_pTransport, -1);
		RT_ASSERTE_RETURN(m_pProxyInfo, -1);

		m_dwSvrIP = aAddr.GetPtr()->sin_addr.s_addr;
		m_wSvrPort = aAddr.GetPtr()->sin_port;

		CRtInetAddr addrProxy(m_pProxyInfo->GetHostName().c_str(), m_pProxyInfo->GetPort());
		return m_TcpConnector.Connect(addrProxy);
	}

	virtual int Close()
	{
		m_TcpConnector.Close();
		if (m_pTransport) {
			m_pTransport = NULL;
		}
		m_pProxyInfo = NULL;
		m_State = STATE_IDLE;
		return 0;
	}

	int OnConnectIndication(
		RtResult aReason, 
		UpTrptType *aTrpt,
		IRtConnectorInternal *aId)
	{
		RT_ASSERTE(m_State == STATE_IDLE);
		RT_ASSERTE(&m_TcpConnector == aId);
		m_pTransport = aTrpt;

		if (RT_SUCCEEDED(aReason)) {
			RT_ASSERTE(m_pTransport);
			aReason = m_pTransport->OpenWithSink(this);
			if (RT_SUCCEEDED(aReason))
				aReason = StartNewRequest();
		}

		if (RT_FAILED(aReason)) {
			Close();
			m_Upper.OnConnectIndication(aReason, NULL, this);
		}
		return 0;
	}

	RtResult StartNewRequest()
	{
		char lpBuf[128];
		DWORD dwBufLen = 0;
		if (m_State == STATE_IDLE) {
			if (m_pProxyInfo->GetProxyType() == CRtHttpProxyInfo::SOCK4_PROXY) {
				lpBuf[0] = 0x04; // SOCKS version 4
				lpBuf[1] = 0x01; // CD command code -- 1 for connect
				::memcpy(lpBuf + 2, &m_wSvrPort, sizeof(m_wSvrPort));
				::memcpy(lpBuf + 4, &m_dwSvrIP, sizeof(m_dwSvrIP));
				::memcpy(lpBuf + 8, "wbx", 3);
				lpBuf[11] = 0x00;
				dwBufLen = 12;
				m_State = STATE_CONNECT_REMOTE;
			}
			else {
				lpBuf[0] = 0x05;// SOCKS version 5
				lpBuf[1] = 0x01; // number of auth procotols we recognize auth protocols
				lpBuf[2] = 0x00; //no authentication required;
				// compliant implementations MUST implement GSSAPI
				// and SHOULD implement username/password and MAY
				// implement CHAP
				// TODO: we don't implement these
				//lpBuf[3] = 0x01; // GSSAPI
				//lpBuf[4] = 0x02; // username/password
				//lpBuf[5] = 0x03; // CHAP
				dwBufLen = lpBuf[1] + 2;
				m_State = STATE_CONNECT_PROXY;
			}
		}
		else if (m_State == STATE_CONNECT_PROXY) {
			RT_ASSERTE(m_pProxyInfo->GetProxyType() == CRtHttpProxyInfo::SOCK5_PROXY);
			lpBuf[0] = 0x05; // SOCKS version 5
			lpBuf[1] = 0x01; // CONNECT command
			lpBuf[2] = 0x00; // obligatory reserved field (perfect for MS tampering!)
			lpBuf[3] = 0x01; // encoding of destination address (1 == IPv4)
			memcpy(lpBuf + 4, &m_dwSvrIP,sizeof(m_dwSvrIP));
			memcpy(lpBuf + 8, &m_wSvrPort, sizeof(m_wSvrPort));
			dwBufLen = 10;
			m_State = STATE_CONNECT_REMOTE;
		}
		else {
			RT_ERROR_TRACE_THIS("CRtConnectorSocksProxyT::StartNewRequest,"
				" wrong state=" << m_State);
			RT_ASSERTE(FALSE);
			return RT_ERROR_UNEXPECTED;
		}

		CRtMessageBlock mbSend(
			dwBufLen, 
			lpBuf, 
			CRtMessageBlock::DONT_DELETE, 
			dwBufLen);
		RtResult rv = m_pTransport->SendData(mbSend);
		if (RT_FAILED(rv)) {
			RT_ERROR_TRACE_THIS("CRtConnectorSocksProxyT::StartNewRequest,"
				" SendData() failed! len=" << dwBufLen << 
				" rv=" << rv);
		}
		return rv;
	}

	virtual void OnReceive(
		CRtMessageBlock &aData,
		IRtTransport *aTrptId,
		CRtTransportParameter *aPara = NULL)
	{
		RT_ASSERTE(!aData.GetNext());
		LPCSTR lpBuf = aData.GetTopLevelReadPtr();
		DWORD dwBufLen = aData.GetTopLevelLength();

		if (m_State == STATE_CONNECT_PROXY) {
			RT_ASSERTE(m_pProxyInfo->GetProxyType() == CRtHttpProxyInfo::SOCK5_PROXY);
			if (dwBufLen >= 2 && (lpBuf[0] == 0x05 && lpBuf[1] == 0x00)) {
				RtResult rv = StartNewRequest();
				if (RT_FAILED(rv))
					goto fail;
				return;
			}
			else {
				RT_ERROR_TRACE_THIS("CRtConnectorSocksProxyT::OnReceive, fail1,"
					" len=" << dwBufLen <<
					" buf[0]=" << lpBuf[0] <<
					" buf[1]=" << lpBuf[1]);
				goto fail;
			}
		}
		else if (m_State == STATE_CONNECT_REMOTE) {
			if (m_pProxyInfo->GetProxyType() == CRtHttpProxyInfo::SOCK4_PROXY) {
				if (dwBufLen >= 8 && (lpBuf[0] == 0x00 && lpBuf[1] == 0x5A))
					m_State = STATE_SUCCESS;
				else {
					RT_ERROR_TRACE_THIS("CRtConnectorSocksProxyT::OnReceive, fail2,"
						" len=" << dwBufLen <<
						" buf[0]=" << lpBuf[0] <<
						" buf[1]=" << lpBuf[1]);
					goto fail;
				}
			}
			else {
				if (dwBufLen >= 10 && (lpBuf[0] == 0x05 && lpBuf[1] == 0x00))
					m_State = STATE_SUCCESS;
				else {
					RT_ERROR_TRACE_THIS("CRtConnectorSocksProxyT::OnReceive, fail3,"
						" len=" << dwBufLen <<
						" buf[0]=" << lpBuf[0] <<
						" buf[1]=" << lpBuf[1]);
					goto fail;
				}
			}
		}

		RT_ASSERTE(m_State == STATE_SUCCESS);
		m_Upper.OnConnectIndication(RT_OK, m_pTransport.ParaIn(), this);
		m_pTransport = NULL;
		return;

fail:
		Close();
		m_Upper.OnConnectIndication(RT_ERROR_NETWORK_SOCKET_ERROR, NULL, this);
		return;
	}

	virtual void OnSend(
		IRtTransport *aTrptId,
		CRtTransportParameter *aPara = NULL)
	{
		RT_ASSERTE(!"CRtConnectorSocksProxyT::OnSend, it shouldn't be called!");
	}

	virtual void OnDisconnect(
		RtResult aReason,
		IRtTransport *aTrptId)
	{
		Close();
		if (RT_SUCCEEDED(aReason))
			aReason = RT_ERROR_NETWORK_SOCKET_ERROR;
		m_Upper.OnConnectIndication(aReason, NULL, this);
	}

private:
	enum SOCKS_STATE 
	{
		STATE_IDLE,
		STATE_CONNECT_PROXY,
		STATE_CONNECT_REMOTE,
		STATE_SUCCESS,
		STATE_FAILURE
	};
	
	IRtReactor *m_pReactor;
	UpperType &m_Upper;
	CRtConnectorTcpT<SelfType, UpTrptType, UpSockType> m_TcpConnector;
	CRtAutoPtr<UpTrptType> m_pTransport;
	SOCKS_STATE m_State;
	CRtAutoPtr<CRtHttpProxyInfo> m_pProxyInfo;

	// follows are in network octet order.
	DWORD m_dwSvrIP;
	WORD  m_wSvrPort;
};

#endif // !RTCONNECTORSOCKSPROXYT_H
